from pymongo import MongoClient
import matplotlib.pyplot as plt
import time

url = 'mongodb://'
client = MongoClient(url)
db = client['2019ncov']


class ncov_analyse:
    pname = ''
    cname = ''
    staType = ''
    sortType = ''
    info = []
    data = []

    def __init__(self):
        self.col = db['raw']

    def findAllByProvince(self):
        items = self.col.find({"provinceShortName": self.pname})
        for i in items:
            self.info.append({
                'confirmedCount': i['confirmedCount'],
                'suspectedCount': i['suspectedCount'],
                'curedCount': i['curedCount'],
                'deadCount': i['deadCount'],
                'updateTime': i['updateTime']
            })
        return self.info

    def findAllByCity(self):
        items = self.col.find({"provinceShortName": self.pname})
        for i in items:
            if ('cities' in i.keys()) and (i['cities'] is not None):
                for j in i['cities']:
                    if j['cityName'] == self.cname:
                        self.info.append({
                            'confirmedCount': j['confirmedCount'],
                            'suspectedCount': j['suspectedCount'],
                            'curedCount': j['curedCount'],
                            'deadCount': j['deadCount'],
                            'updateTime': i['updateTime']
                        })
        return self.info

    def sortInfoByDay(self, dataNum, dataTime):
        newNum = []
        newTime = []
        length = len(dataTime)
        for index, value in enumerate(dataTime):
            if value != dataTime[index + 1]:
                newNum.append(dataNum[index])
                newTime.append(dataTime[index])
            if index == (length - 2):
                break
        if newTime[len(newTime) - 1] != dataTime[length - 1]:
            newNum.append(dataNum[length - 1])
            newTime.append(dataTime[length - 1])
        return [newNum, newTime]

    def formatInfo(self):
        dataNum = []
        dataTime = []
        for i in self.info:
            dataNum.insert(0, i[self.sortType])
            dataTime.insert(
                0,
                time.strftime("%m-%d",
                              time.localtime(int(i['updateTime'] / 1000))))
        return self.sortInfoByDay(dataNum, dataTime)

    def calTotal(self):
        self.data = self.formatInfo()
        return self.data

    def calNewlyIncreased(self):
        tmpData = self.formatInfo()
        tmpNum = []
        length = len(tmpData[0])
        for index, value in enumerate(tmpData[0]):
            tmpNum.append(tmpData[0][index + 1] - value)
            if index == (length - 2):
                break
        tmpData[1].pop()
        self.data = [tmpNum, tmpData[1]]
        return self.data

    def calDeathRate(self):
        pass

    def calCureRate(self):
        pass

    def draw(self):
        plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']
        plt.rcParams['axes.unicode_minus'] = False
        plt.plot(self.data[1],
                 self.data[0],
                 label=self.sortType,
                 linewidth=3,
                 color='r',
                 marker='o',
                 markerfacecolor='blue',
                 markersize=12)
        plt.xlabel('date')
        plt.ylabel('value')
        plt.title('2019n-cov\n' + self.staType + '-' + self.sortType + '-' +
                  self.pname + self.cname)
        for a, b in zip(self.data[1], self.data[0]):
            plt.text(a, b, b, ha='center', va='bottom', fontsize=10)
        plt.legend()
        plt.show()
